# -*- coding: binary -*-

module Msf
  module Exploit::Remote::SMB::Server
    # This mixin provides a minimal SMB server sharing an UNC resource. At
    # this moment it is capable to share just one file. And the file should
    # live in the root folder "\\".
    module Share
      include ::Msf::Exploit::Remote::SMB::Server
      include ::Msf::Exploit::Remote::SMB::Server::HashCapture

      # @!attribute share
      #   @return [String] The share portion of the provided UNC.
      attr_accessor :share
      # @!attribute folder_name
      #   @return [String] The folder where the provided file lives.
      attr_accessor :folder_name
      # @!attribute file_name
      #   @return [String] The file name of the provided UNC.
      attr_accessor :file_name
      # @!attribute file_contents
      #   @return [String] The contents of the provided file
      attr_accessor :file_contents

      def initialize(info = {})
        super

        register_options(
          [
            OptString.new('SHARE', [ false, 'Share (Default: random); cannot contain spaces or slashes'], regex: /^[^\s\/\\]*$/),
            OptString.new('FILE_NAME', [ false, 'File name to share (Default: random)']),
            OptString.new('FOLDER_NAME', [ false, 'Folder name to share (Default: none)'])
          ], Msf::Exploit::Remote::SMB::Server::Share)
        register_advanced_options(
          [
            OptString.new('SMBDomain', [ true, 'The domain name used during SMB exchange.', 'WORKGROUP'])
          ]
        )
      end

      def start_service(opts = {})
        unless opts[:gss_provider]
          ntlm_provider = Msf::Exploit::Remote::SMB::Server::HashCapture::HashCaptureNTLMProvider.new(
            allow_anonymous: true,
            allow_guests: true,
            listener: self,
            ntlm_type3_status: nil
          )

          # Set domain name for all future server responses
          ntlm_provider.dns_domain = datastore['SMBDomain']
          ntlm_provider.dns_hostname = datastore['SMBDomain']
          ntlm_provider.netbios_domain = datastore['SMBDomain']
          ntlm_provider.netbios_hostname = datastore['SMBDomain']
          opts[:gss_provider] = ntlm_provider
        end

        super(opts)

        if share.present?
          if service.shares.key?(share)
            fail_with(Msf::Module::Failure::BadConfig, "The specified SMB share '#{share}' already exists.")
          end

          virtual_disk = RubySMB::Server::Share::Provider::VirtualDisk.new(share)
          # the virtual disk expects the path to use the native File::SEPARATOR so normalize on that here
          virtual_disk.add_dynamic_file("#{@folder_name}#{File::SEPARATOR}#{@file_name}".gsub(/\/|\\/, File::SEPARATOR)) do |client, _smb_session|
            get_file_contents(client: client)
          end
          service.add_share(virtual_disk)
        end
      end

      # Setups the server configuration.
      def setup
        super

        self.folder_name = datastore['FOLDER_NAME']
        self.share = datastore['SHARE'].present? ? datastore['SHARE'] : Rex::Text.rand_text_alpha(4 + rand(3))
        self.file_name = datastore['FILE_NAME'].present? ? datastore['FILE_NAME'] : Rex::Text.rand_text_alpha(4 + rand(3))
      end

      # Builds the UNC Name for the shared file
      def unc
        if folder_name
          path = "\\\\#{srvhost}\\#{share}\\#{folder_name}\\#{file_name}"
        else
          path = "\\\\#{srvhost}\\#{share}\\#{file_name}"
        end

        path
      end

      # Returns the file contents for the requested file
      #
      # @return [String] The file contents.
      def get_file_contents(client:)
        file_contents
      end

      def cleanup
        self.service.remove_share(share) if share.present?

        super
      end
    end
  end
end
